#include <iostream>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <chrono>
#include <thread>
#include <atomic>

#include <unitree/common/thread/thread.hpp>
#include <unitree/common/time/time_tool.hpp>
#include <unitree/idl/go2/LowState_.hpp>
#include <unitree/robot/channel/channel_subscriber.hpp>
#include <unitree/robot/go2/sport/sport_client.hpp>

using namespace unitree::robot;
using namespace unitree::common;
using namespace std;

#define TOPIC_LOWSTATE "rt/lowstate"

// 获取 yaw 角度（弧度），范围 [-π, π]
float GetYaw(const unitree_go::msg::dds_::LowState_& state) {
    return state.imu_state().rpy()[2];  // yaw 是第三个元素
}

// 计算角度差（考虑 wrap-around），范围 [0, π]
float YawDiff(float current, float target) {
    float diff = fabs(current - target);
    if (diff > M_PI) diff = 2 * M_PI - diff;
    return diff;
}

int main(int argc, const char** argv)
{
    ChannelFactory::Instance()->Init(0, "eth0");
    cout << "Factory Init" << endl;

    // 初始化 IMU 订阅
    unitree_go::msg::dds_::LowState_ low_state;
    atomic<float> current_yaw {0.0f};

    auto sub = std::make_shared<ChannelSubscriber<unitree_go::msg::dds_::LowState_>>(TOPIC_LOWSTATE);
    sub->InitChannel([&](const void* msg) {
        low_state = *(unitree_go::msg::dds_::LowState_*)msg;
        current_yaw.store(GetYaw(low_state));
    }, 1);

    // 等待 IMU 数据稳定
    this_thread::sleep_for(chrono::seconds(1));

    // 记录初始 yaw
    float start_yaw = current_yaw.load();
    float target_yaw = start_yaw + M_PI_2;  // 加 90 度（π/2 弧度）
    if (target_yaw > M_PI) target_yaw -= 2 * M_PI;

    cout << "Start Yaw: " << start_yaw << ", Target Yaw: " << target_yaw << endl;

    // 初始化运动控制
    unitree::robot::go2::SportClient spt;
    spt.SetTimeout(10.0f);
    spt.Init();

    float vx = 0.0f, vy = 0.0f, vyaw = 0.7f;  // 角速度约 0.7 rad/s

    // 控制循环
    while (true) {
        float cyaw = current_yaw.load();
        float diff = YawDiff(cyaw, target_yaw);

        cout << "Current Yaw: " << cyaw << ", Diff: " << diff << endl;

        if (diff < 0.05f) {
            cout << "Target reached!" << endl;
            break;
        }

        spt.Move(vx, vy, vyaw);
        this_thread::sleep_for(chrono::milliseconds(100));
    }

    spt.StopMove();
    cout << "Rotation done. Robot stopped." << endl;

    return 0;
}
